import h5py
import numpy as np
import tensorflow as tf
from .config import *
from typing import Generator, Tuple

def day_block_shuffle(times):
    """按天分块shuffle"""
    dates = np.array([t.date() for t in times])
    unique_dates = np.unique(dates)
    blocks = [np.where(dates == d)[0] for d in unique_dates]
    np.random.shuffle(blocks)
    return np.concatenate(blocks)

def process_images(images):
    """图像预处理"""
    images = tf.transpose(images, [0, 2, 3, 1, 4])
    images = tf.reshape(images, [-1, config.IMG_SIDE_LEN, config.IMG_SIDE_LEN, config.NUM_LOG_TERM*config.NUM_COLOR_CHANNEL])
    return tf.image.convert_image_dtype(images, tf.float32)

def create_dataset(data_path: str, indices: list, batch_size: int) -> tf.data.Dataset:
    """创建支持大文件的TF Dataset"""
    # 确保索引有序且去重
    indices = sorted(list(set(indices)))
    total_samples = len(indices)
    
    # 定义流式数据生成器
    def data_generator() -> Generator:
        # 分块参数设置（根据可用内存调整）
        chunk_size = min(2048, total_samples)  # 每次加载的样本数
        
        with h5py.File(data_path, 'r') as f:
            image_ds = f['trainval']['image_log_trainval']
            pv_log_ds = f['trainval']['pv_log_trainval']
            pv_pred_ds = f['trainval']['pv_pred_trainval']
            
            # 分块处理避免全量加载
            for chunk_start in range(0, total_samples, chunk_size):
                chunk_end = min(chunk_start + chunk_size, total_samples)
                chunk_indices = indices[chunk_start:chunk_end]
                
                # 按需加载当前块数据
                images = process_images(image_ds[sorted(chunk_indices)])
                pv_log = pv_log_ds[sorted(chunk_indices)]
                pv_pred = pv_pred_ds[sorted(chunk_indices)]
                
                # 确保块内顺序
                order = np.argsort(chunk_indices)
                for i in order:
                    if config.MODEL_SELECT == "CNN_LSTM":
                        yield (images[i], pv_log[i]), pv_pred[i]
                    else:
                        yield pv_log[i], pv_pred[i]

    # 定义输出签名
    if config.MODEL_SELECT == "CNN_LSTM":
        output_signature = (
            (tf.TensorSpec(shape=(64,64,48), dtype=tf.float32),
             tf.TensorSpec(shape=(16,), dtype=tf.float32)),
            tf.TensorSpec(shape=(), dtype=tf.float32)
        )
    else:
        output_signature = (
            tf.TensorSpec(shape=(16,), dtype=tf.float32),
            tf.TensorSpec(shape=(), dtype=tf.float32)
        )

    # 构建数据集
    dataset = tf.data.Dataset.from_generator(
        data_generator,
        output_signature=output_signature
    )
    
    # 优化配置
    return dataset\
        .shuffle(len(indices))\ #随机打乱数据，增强模型泛化性
        .batch(batch_size)\
        .prefetch(tf.data.AUTOTUNE)\
        .cache()  # 可选缓存最近使用的数据块